package org.shj.spark.udaf

import org.apache.spark.sql.expressions.UserDefinedAggregateFunction
import org.apache.spark.sql.types._
import org.apache.spark.sql.expressions.MutableAggregationBuffer
import org.apache.spark.sql.Row
import org.apache.spark.sql.SparkSession
import org.shj.spark.util.Util

/**
 * 本例演示 untyped 自定义聚合函数
 */
object MyAverageUDF extends UserDefinedAggregateFunction{
  
  //自定义的聚合函数的输入参数的数据类型
  def inputSchema: StructType = StructType(StructField("inputColumn", LongType) :: Nil)
  
  // 在进行聚合操作的时候所要处理的数据的结果的类型
  def bufferSchema: StructType = {
    StructType(StructField("sum", LongType)::StructField("count", LongType) :: Nil)
  }
  
  // 最终返回的结果的类型
  def dataType: DataType = DoubleType
  
  //相同的输入是否总是得到相同的输出。一般情况下都为true
  def deterministic: Boolean = true
  
  //初始化给定的 aggregation buffer。 buffer本身是 "Row"。它除了包含标准的方法(例如： get(), getBoolean())外，还
  //提供了更新值的机会。 注意： buffer 中的 Array 和 Map仍然是不可变的。 
  def initialize(buffer: MutableAggregationBuffer): Unit = {
    buffer(0) = 0L
    buffer(1) = 0L
  }
  
  // Updates the given aggregation buffer `buffer` with new input data from `input`
  //在进行聚合的时候，每当有新的值进来，对分组后的聚合如何进行计算
  //本地的聚合操作，相当于MapReduce模型中的combiner
  //用 "input" 去更新 buffer
  def update(buffer: MutableAggregationBuffer, input: Row)={
    if(!input.isNullAt(0)){
      buffer(0) = buffer.getLong(0) + input.getLong(0)
      buffer(1) = buffer.getLong(1) + 1
    }
  }
  
  // Merges two aggregation buffers and stores the updated buffer values back to `buffer1`
  // 最后在分布式节点进行Local Reduce完成后需要进行全局级别的Merge操作
  //合并两个buffer，并把结果更新到 buffer1
  def merge(buffer1: MutableAggregationBuffer, buffer2: Row)={
    buffer1(0) = buffer1.getLong(0) + buffer2.getLong(0)
    buffer1(1) = buffer1.getLong(1) + buffer2.getLong(1)
  }
  
  //计算最后结果
  def evaluate(buffer: Row) : Double = buffer.getLong(0).toDouble / buffer.getLong(1)
  
  def main(args: Array[String]): Unit = {
    val spark = SparkSession.builder().appName("MyAverageUDF").master("local").getOrCreate()
    spark.sparkContext.setLogLevel("WARN")
    
    //注册自定义的函数
    spark.udf.register("myAverage", MyAverageUDF)
    
    val df = spark.read.json(Util.fullPath("employees.json"))
    df.createOrReplaceTempView("emp")
    df.show()
    
    spark.sql("select deptCd, myAverage(salary) as avg_salary from emp group by deptCd").show()
    
    spark.stop()
  }
}